package flowdroid.taint;

import flowdroid.scan.MethodChecker;
import pojo.MyTuple;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.jimple.infoflow.android.data.AndroidMethod;
import soot.jimple.infoflow.collect.ConcurrentHashSet;
import soot.jimple.infoflow.sourcesSinks.definitions.ISourceSinkDefinition;
import soot.jimple.infoflow.sourcesSinks.definitions.ISourceSinkDefinitionProvider;
import soot.jimple.infoflow.sourcesSinks.definitions.MethodSourceSinkDefinition;
import soot.jimple.infoflow.sourcesSinks.definitions.SourceSinkType;
import tuple.data.DataKind;
import util.Constant;
import util.FlowdroidUtils;

import java.io.FileWriter;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * My Custom Source Sink provider
 */
public class MyProvider implements ISourceSinkDefinitionProvider {
    Set<ISourceSinkDefinition> sourceList;
    Set<ISourceSinkDefinition> sinkList;
    Set<ISourceSinkDefinition> methods;

    public MyProvider() {
        this.sourceList = new HashSet<>();
        this.sinkList = new HashSet<>();
        this.methods = new HashSet<>();
    }

    public void clear() {
        this.sourceList.clear();
        this.sinkList.clear();
        this.methods.clear();
    }


    @Override
    public Collection<? extends ISourceSinkDefinition> getSources() {
        return this.sourceList;
    }

    @Override
    public Collection<? extends ISourceSinkDefinition> getSinks() {
        return this.sinkList;
    }

    @Override
    public Collection<? extends ISourceSinkDefinition> getAllMethods() {
        return this.methods;
    }

    public void addAll(ISourceSinkDefinitionProvider other) {
        this.sourceList.addAll(other.getSources());
        this.sinkList.addAll(other.getSinks());
        this.methods.addAll(other.getAllMethods());
    }

    public void addSourceFromSootMethod(Collection<SootMethod> target) {
        for (SootMethod s : target) {
            AndroidMethod androidMethod = new AndroidMethod(s);
            androidMethod.setSourceSinkType(SourceSinkType.Source);
            this.sourceList.add(new MethodSourceSinkDefinition(androidMethod));
        }
    }

    public void addSourceFromDefinitionSet(Set<ISourceSinkDefinition> target) {
        try {
            this.sourceList.addAll(target);
        } catch (Exception e) {
            System.out.println("Error in addSourceByList(Set<ISourceSinkDefinition> target) :" + e);
        }
    }

    public void addSinkFromSootMethod(Collection<SootMethod> target) {
        for (SootMethod s : target) {
            AndroidMethod androidMethod = new AndroidMethod(s);
            androidMethod.setSourceSinkType(SourceSinkType.Sink);
            this.sinkList.add(new MethodSourceSinkDefinition(androidMethod));
        }
    }

    public void addSinkFromDefinitionSet(Set<ISourceSinkDefinition> target) {
        try {
            this.sinkList.addAll(target);
        } catch (Exception e) {
            System.out.println("Error in addSinkByList(Set<ISourceSinkDefinition> target) :" + e);
        }
    }

    public void removeSourceFromSootMethod(Collection<SootMethod> target) {
        for (SootMethod s : target) {
            AndroidMethod androidMethod = new AndroidMethod(s);
            androidMethod.setSourceSinkType(SourceSinkType.Source);
            for (ISourceSinkDefinition idef : this.sourceList) {
                if (idef instanceof AndroidMethod am) {
                    if (am.getSignature().equals(androidMethod.getSignature())) {
                        this.sourceList.remove(idef);
                    }
                }
            }
        }
    }

    private void removeByString(String target, Set<ISourceSinkDefinition> list) {
        ConcurrentHashSet<ISourceSinkDefinition> toRemove = new ConcurrentHashSet<>();
        for (ISourceSinkDefinition idef : list) {
            if (idef instanceof MethodSourceSinkDefinition mdef) {
                if (mdef.getMethod().getSignature().contains(target)) {
                    toRemove.add(idef);
                }
            }
        }
        list.removeAll(toRemove);
    }

    public void removeSourceByString(String target) {
        removeByString(target, this.sourceList);
    }

    public void removeSinkByString(String target) {
        removeByString(target, this.sinkList);
    }


    public void syncAfterChange() {
        HashSet<MethodSourceSinkDefinition> newMethods = new HashSet<>();
        newMethods.addAll(this.sourceList.stream()
                .filter(m -> m instanceof MethodSourceSinkDefinition)
                .map(m -> (MethodSourceSinkDefinition) (m)).collect(Collectors.toSet()));
        newMethods.addAll(this.sinkList.stream()
                .filter(m -> m instanceof MethodSourceSinkDefinition)
                .map(m -> (MethodSourceSinkDefinition) (m)).collect(Collectors.toSet()));
        if (methods.size() != newMethods.size() || !methods.equals(newMethods)) {
            this.methods.clear();
            this.methods.addAll(newMethods);
        }
    }


    public void dumpSourceSinkToFile(String filepath) {
        syncAfterChange();
        System.out.printf("Write Sources: %d, Sinks: %d To file: <%s> \n", sourceList.size(), sinkList.size(), filepath);

        try (FileWriter fw = new FileWriter(filepath)) {
            sourceList.forEach(
                    x -> {
                        try {
                            fw.write(x.toString() + " -> _SOURCE_\n");
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
            );

            sinkList.forEach(
                    x -> {
                        try {
                            fw.write(x.toString() + " -> _SINK_\n");
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
            );
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void loadMarkedCoreMethodsAsSource(MethodChecker marker) {
        Set<SootMethod> coreMethods = marker.getCoreMethods();
        Map<SootMethod, MyTuple> methodsMap = marker.methodsMap;
        if (coreMethods.size() < Constant.CORE_METHODS_SIZE_BOUND) {
            this.addSourceFromSootMethod(coreMethods);
        } else {
            HashSet<SootMethod> toAdd = new HashSet<>();
            HashSet<DataKind> dataKinds = new HashSet<>();
            for (SootMethod sm : coreMethods) {
                Set<DataKind> smData;
                try {
                    smData = methodsMap.get(sm).data;
                } catch (Exception e) {
                    smData = new HashSet<>();
                }

                if (toAdd.size() < Constant.CORE_METHODS_SIZE_BOUND) {
                    toAdd.add(sm);
                    if (smData != null)
                        dataKinds.addAll(smData);
                } else if (smData.stream().anyMatch(x -> !dataKinds.contains(x))) {
                    toAdd.add(sm);
                    dataKinds.addAll(smData);
                }
            }
            this.addSourceFromSootMethod(toAdd);
        }

        Arrays.stream(FlowdroidUtils.excludeSourceMethodList).forEach(
                this::removeSourceByString
        );

//        this.addSinkByList(addSinks.stream().toList());
//        Arrays.stream(Constant.excludeSinkMethodList).forEach(
//                this::removeSinkByString
//        );

        syncAfterChange();
    }

    @Deprecated
    public void loadMoreSourceSink() {
        List<SootMethod> target = new ArrayList<>();
        // android.os.Build 里 很多字段 Field 也应该用污点
        Arrays.stream(FlowdroidUtils.sentitiveSources).forEach(str -> {
            Scene.v().loadClass(str, SootClass.BODIES);
            List<SootMethod> ta = Scene.v().getSootClass(str).getMethods();
            target.addAll(ta.stream().filter(x -> x.getName().contains("get") || x.getName().contains("query")
                            || x.getName().contains("request") || x.getName().contains("access")
                            || x.getName().contains("fetch") || x.getName().contains("obtain")
                            || x.getName().contains("acquire")
                    )
                    .toList());
        });


        this.addSourceFromSootMethod(target);

        Arrays.stream(FlowdroidUtils.excludeSourceMethodList).forEach(
                this::removeSourceByString
        );

        List<SootMethod> targetSinks = new ArrayList<>();
        Arrays.stream(FlowdroidUtils.sentitiveSinks).forEach(str -> {
            Scene.v().loadClass(str, SootClass.BODIES);
            List<SootMethod> ta = Scene.v().getSootClass(str).getMethods();
            targetSinks.addAll(ta.stream().filter(x -> x.getName().contains("put") || x.getName().contains("set")
                            || x.getName().contains("send") || x.getName().contains("dump") || x.getName().contains("post"))
                    .toList());
        });
        this.addSinkFromSootMethod(targetSinks);

        Arrays.stream(FlowdroidUtils.excludeSinkMethodList).forEach(
                this::removeSinkByString
        );
        syncAfterChange();
    }
//
//    @Deprecated
//    public void loadPermissionFlowSources(MethodChecker marker) {
//        this.addSourceFromSootMethod(marker.getPermissionMethods());
//        this.addSinkFromSootMethod(marker.getPermissionMethods());
//
//        Arrays.stream(FlowdroidUtils.excludeSourceMethodList).forEach(
//                this::removeSourceByString
//        );
//
////        List<SootMethod> targetSinks = new ArrayList<>();
////        Arrays.stream(Constant.sentitiveSinks).forEach(str -> {
////            Scene.v().loadClass(str, SootClass.BODIES);
////            List<SootMethod> ta = Scene.v().getSootClass(str).getMethods();
////            targetSinks.addAll(ta.stream().filter(x -> x.getName().contains("put") || x.getName().contains("set")
////                            || x.getName().contains("send") || x.getName().contains("dump") || x.getName().contains("post"))
////                    .collect(Collectors.toList()));
////        });
////        this.addSinkByList(targetSinks);
////
////        Arrays.stream(Constant.excludeSinkMethodList).forEach(
////                this::removeSinkByString
////        );
//        syncAfterChange();
//    }

    @Deprecated
    public static Set<SootMethod> findPermissionSources() {
        Set<SootMethod> target = new HashSet<>();
        Set<SootClass> applicationClazz = Scene.v().getClasses().stream()
                .filter(cl ->
                        !(cl.getName().startsWith("java") || cl.getName().startsWith("android")))
                .collect(Collectors.toSet());

        for (SootClass sootClass : applicationClazz) {
//            sootClass.setApplicationClass();
            Set<SootMethod> hasBodyMethods = sootClass.getMethods().stream()
                    .filter(SootMethod::hasActiveBody)
                    .collect(Collectors.toSet());

            Set<SootMethod> containsPermissionMethods = hasBodyMethods.stream()
                    .filter(m ->
                            m.retrieveActiveBody()
                                    .getUnits()
                                    .stream()
                                    .anyMatch(u -> u.toString().contains("android.permission.")))
                    .collect(Collectors.toSet());

            target.addAll(containsPermissionMethods);
        }
        return target;
//        this.addSourceByList(target);
//
//        Arrays.stream(Constant.excludeSourceMethodList).forEach(
//                this::removeSourceByString
//        );
//        syncMethods();
    }


}